#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File     : loader.py
@Project  : pipecoco
@Date     : 2021/8/23
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''

from mindspore import load_checkpoint, load_param_into_net
import os
from net.alexnet.alexnet import AlexNet
from pipecoco.utils import dfs_search_layer
import logging

logger = logging.getLogger("net")

def load_alexnet(model_path = None):

    network = AlexNet(num_classes=10)
    if model_path!=None and os.path.exists(model_path):
        param_dict = load_checkpoint(model_path)
        load_param_into_net(network, param_dict)
        logger.info("load checkpoint from [{}].".format(model_path))
        logger.info("{} load checkpoint from [{}].".format(network.cls_name,model_path))
    network.set_train(False)

    cell_names = ['conv1', 'relu', 'max_pool2d', 'conv2', 'relu', 'max_pool2d','conv3', 'relu','pad', 'conv4', 'relu', 'conv5', 'relu', 'max_pool2d','flatten', 'fc1', 'relu', 'fc2', 'relu', 'fc3']
    cells_list = []
    for name in cell_names:
        cells_list.append(getattr(network, name))
    layers_list = []
    for cell in cells_list:
        dfs_search_layer(cell, layers_list)
    return layers_list